#include <cuda_runtime.h>
#include <cusparse.h>
#include <stdio.h>
#include <stdlib.h>
#include "toroidal.h"

/* cuda memory transfer */

void paragrad(int* offs, int* col, double* vals, int mdim, double* br,
double* bz, double* bp)
{
    double ivdr = 1.0/dr;
    double ivdz = 1.0/dz;
    int nmode = 4; // toroidal mode number
    for(int i = 0; i < mdim/2; i++)
    {
        if (i==0)
            offs[i] = 0;
        offs[i+1] = offs[i] + 5;
        int ix = i/Ny;
        int iy = i%Ny;
        int dn = (iy == 0 ? i : i-1);
        int up = (iy == Ny-1 ? i : i+1);
        int lf = (ix == 0 ? i : i-Ny);
        int rg = (ix == Nx-1 ? i : i+Ny);
        col[5*i] = up;
        col[5*i+1] = dn;
        col[5*i+2] = lf;
        col[5*i+3] = rg;
        vals[5*i] = -br[i]*ivdr/(double)(up-dn);
        vals[5*i+1] = br[i]*ivdr/(double)(up-dn);
        vals[5*i+2] = -bz[i]*ivdz/(double)((rg-lf)/Ny);
        vals[5*i+3] = bz[i]*ivdz/(double)((rg-lf)/Ny);
        col[5*i+4] = i+mdim/2;
        vals[5*i+4] = -bp[i]*nmode;
    }
    for(int iq = mdim/2; iq < mdim; iq++)
    {
        int i = iq - mdim/2;
        offs[i+1] = offs[i] + 5;
        int ix = i/Ny;
        int iy = i%Ny;
        int dn = (iy == 0 ? i : i-1);
        int up = (iy == Ny-1 ? i : i+1);
        int lf = (ix == 0 ? i : i-Ny);
        int rg = (ix == Nx-1 ? i : i+Ny);
        col[5*iq] = up + mdim/2;
        col[5*iq+1] = dn + mdim/2;
        col[5*iq+2] = lf + mdim/2;
        col[5*iq+3] = rg + mdim/2;
        vals[5*iq] = -br[i]*ivdr/(double)(up-dn);
        vals[5*iq+1] = br[i]*ivdr/(double)(up-dn);
        vals[5*iq+2] = -bz[i]*ivdz/(double)((rg-lf)/Ny);
        vals[5*iq+3] = bz[i]*ivdz/(double)((rg-lf)/Ny);
        col[5*iq+4] = iq-mdim/2;
        vals[5*iq+4] = bp[i]*nmode;
    }
}

void laplacian(int* offs, int* col, double* vals, int mdim)
{
    double ivdr = 1.0/dr;
    double ivdz = 1.0/dz;
    for(int i = 0; i < mdim; i++)
    {
        if (i==0)
            offs[i] = 0;
        offs[i+1] = offs[i] + 5;
        int ix = i/Ny;
        int iy = i%Ny;
        if(ix == 0) {
            col[5*i+2] = i+2*Ny;
            vals[5*i+2] = ivdr*ivdr;
            col[5*i+1] = i+Ny;
            vals[5*i+1] = -2.0*ivdr*ivdr;
            col[5*i] = i;
            vals[5*i] = ivdr*ivdr;
        }
        else if (ix == Nx-1)
        {
            col[5*i+2] = i-2*Ny;
            vals[5*i+2] = -ivdr*ivdr;
            col[5*i+1] = i-Ny;
            vals[5*i+1] = 2.0*ivdr*ivdr;
            col[5*i] = i;
            vals[5*i] = -ivdr*ivdr;
        }
        else 
        {
            col[5*i] = i-Ny;
            vals[5*i] = ivdr*ivdr;
            col[5*i+1] = i;
            vals[5*i+1] = -2.0*ivdr*ivdr;
            col[5*i+2] = i+Ny;
            vals[5*i+2] = -ivdr*ivdr;
        }
        if(iy == 0) {
            vals[5*i] += ivdz*ivdz;
            col[5*i+3] = i+2;
            vals[5*i+3] = ivdz*ivdz;
            col[5*i+4] = i+1;
            vals[5*i+4] = -2.0*ivdz*ivdz;
            
        }
        else if (iy == Ny-1)
        {
            col[5*i+3] = i-2;
            vals[5*i+3] = -ivdz*ivdz;
            col[5*i+4] = i-1;
            vals[5*i+4] = 2.0*ivdz*ivdz;
            vals[5*i] += -ivdz*ivdz;
        }
        else 
        {
            col[5*i+3] = i-1;
            vals[5*i+3] = ivdz*ivdz;
            vals[5*i] += -2.0*ivdz*ivdz;
            col[5*i+4] = i+1;
            vals[5*i+4] = -ivdz*ivdz;
        }
        
    }
}

void cudaMvk(cusparseDnVecDescr_t xvec, cusparseDnVecDescr_t bvec, 
cusparseHandle_t handle, cusparseSpMatDescr_t mat,int mdim) // perform Ax = b
{
    double alpha = 1.0f;
    double beta = 0.0f;

    void* dbuf = NULL;
    size_t bufsize = 0;
    
    cusparseSpMV_bufferSize(handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
    &alpha, mat, xvec, &beta, bvec, CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, &bufsize);    
    cudaMalloc(&dbuf, bufsize);
    cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
    &alpha, mat, xvec, &beta, bvec, CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, dbuf);

    return;
}

cusparseSpMatDescr_t cuMat(int* offsets, int* columns, double* values, int mdim, int nval,
cusparseHandle_t handle) //nval: number of elements
{   
    int *dcsr_offsets, *dcsr_columns;
    double* dcsr_values;
    cudaMalloc((void**)&dcsr_offsets, (mdim+1)*sizeof(int));
    cudaMalloc((void**)&dcsr_columns, nval*sizeof(int));
    cudaMalloc((void**)&dcsr_values, nval*sizeof(double));

    cudaMemcpy(dcsr_offsets, offsets, (mdim+1)*sizeof(int),cudaMemcpyHostToDevice);
    cudaMemcpy(dcsr_columns, columns, nval*sizeof(int),cudaMemcpyHostToDevice);
    cudaMemcpy(dcsr_values, values, nval*sizeof(double),cudaMemcpyHostToDevice);

    cusparseSpMatDescr_t mat;
    cusparseCreate(&handle);
    cusparseCreateCsr(&mat, mdim, mdim, nval, dcsr_offsets, dcsr_columns,dcsr_values,
        CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_64F); 
    
    return mat;
}

void cutest(double* input, double* ret, int mdim)
{
    int nval = 5*mdim;
    int csr_offsets[mdim+1] = {0};    
    int csr_columns[nval] = {0};
    double csr_values[nval] = {0.0};
    laplacian(csr_offsets, csr_columns, csr_values, Nx*Ny);
   // double* ret = new double[Nx*Ny];
    
    double *dx, *db; 
    cudaMalloc((void**)&dx, mdim*sizeof(double));
    cudaMalloc((void**)&db, mdim*sizeof(double));
    cudaMemcpy(dx, input, mdim*sizeof(double),cudaMemcpyHostToDevice);
    cusparseHandle_t handle;
    cusparseDnVecDescr_t xvec;
    cusparseDnVecDescr_t bvec;
    cusparseCreate(&handle);
    cusparseCreateDnVec(&xvec, mdim, dx, CUDA_R_64F);
    cusparseCreateDnVec(&bvec, mdim, db, CUDA_R_64F);
    cusparseSpMatDescr_t mat = cuMat(csr_offsets, csr_columns, csr_values,
     mdim, nval, handle);
    cudaMvk(xvec, bvec, handle, mat, Nx*Ny);
    cudaMemcpy(ret, db, mdim*sizeof(double),cudaMemcpyDeviceToHost);

    cudaDeviceSynchronize();  // very important !!!

}
